Pytorch Shift Stack Demo¶

We demonstrate the shift and stack

Loading Data + Preprocessing¶

We first loop through the directories of where the data lives after quering the data

In [1]:
import os
directory = os.fsencode('Data/aligned')
files = []
for sub1 in os.listdir(directory):
    sub1_name = os.fsdecode(sub1)
    if '.fits'in sub1_name:
        files.append('Data/aligned'+"/"+sub1_name)

We can then load the fits and keep track of the times in a dictionary. NOTE: These files are preprocessed and aligned for RA/DEC offsets

In [2]:
import os
from tqdm import tqdm
from astropy.io import fits
from astropy.time import TimeDelta, Time
import torch
import numpy as np
from astropy.io import fits
import matplotlib.pyplot as plt
from astropy.visualization import (ImageNormalize, ZScaleInterval, LogStretch, MinMaxInterval, AsinhStretch)
from astropy.wcs import WCS
from scipy.ndimage import shift
from astropy.coordinates import SkyCoord

total_data = {}
pass_counter = 0
MAGZPL_LIST = []
for i in tqdm(range(len(files))):

    # Open the FITS file
    filename = files[i]
    hdul = fits.open(filename)    
    # Access the primary HDU (Header Data Unit)
    primary_hdu = hdul[0]
    
    # Get the data and header
    # try:
    data = primary_hdu.data
    header = primary_hdu.header

    total_elements = data.size
    # Count the number of NaN values
    nan_count = np.isnan(data).sum()
    # Calculate the percentage of NaN values
    nan_percentage = (nan_count / total_elements) * 100
    if nan_percentage < 20:
        time = header['SHUTOPEN']
        MAGZPL_LIST.append(header['MAGZP'])
        time = Time(time.replace('T',' '), format='iso', scale='utc')
        total_data[time] = torch.tensor(data.astype(np.float32))
    else:
        pass_counter += 1
        pass
    hdul.close()
print(pass_counter, 'skipped out of ', len(files)) 
MAGZPL = np.max(MAGZPL_LIST)
print(MAGZPL)
100%|█████████████████████████████████████████| 939/939 [00:16<00:00, 57.27it/s]
671 skipped out of  939
26.416529

Get the earliest time of observation

In [3]:
sorted_keys = sorted(total_data.keys())
earliest_frame = sorted_keys[0]

Load the earliest data frame and plot what it looks like

In [4]:
# Apply normalization
norm = ImageNormalize(stretch=AsinhStretch(a=0.0001), interval=ZScaleInterval())  
data = total_data[earliest_frame]
# norm = ImageNormalize(total_data[earliest_frame], interval=ZScaleInterval())

# Plot the image with normalization
plt.figure(figsize=(10, 8))
plt.imshow(data, cmap='gray_r', origin='lower', norm=norm)
plt.colorbar(label='Pixel Value')
plt.title('FITS Image with Normalization')
plt.xlabel('RA Pixel')
plt.ylabel('Dec Pixel')
plt.show()
No description has been provided for this image

To verify we have aligned the frames together we plot a movie of the frames

In [5]:
from astropy.io import fits
import matplotlib.pyplot as plt
from astropy.visualization import (ImageNormalize, ZScaleInterval, LogStretch, MinMaxInterval, AsinhStretch)
import numpy as np
import matplotlib.pyplot as plt
from celluloid import Camera
from IPython.display import HTML

norm = ImageNormalize(stretch=AsinhStretch(a=0.0001), interval=ZScaleInterval())  

%matplotlib inline
# Set up the figure and axes
fig, ax = plt.subplots(figsize=(8,8))

# Initialize Celluloid Camera
camera = Camera(fig)


data = total_data[earliest_frame]
count = 0
# Animate the sine wave with a changing phase
for key in sorted_keys:
    ax.imshow(total_data[key], cmap='gray_r', origin='lower', norm=norm)  # Plot the sine wave
    camera.snap()  # Capture the frame
    if count > 100:
        break
    else:
        count += 1

# Create the animation
animation = camera.animate(interval=100)  # 100ms between frames

# Display the animation
HTML(animation.to_jshtml())  # Render as JS HTML
Out[5]:
No description has been provided for this image
No description has been provided for this image

Add fakes¶

We add gaussian dot that shifts across the sky as a function of time.

In [6]:
import numpy as np
import matplotlib.pyplot as plt

def generate_2d_gaussian(grid_size, amplitude=1, x0=0, y0=0, sigma_x=1, sigma_y=1):
    Nx, Ny = grid_size
    x = np.linspace(-Nx // 2, Nx // 2, Nx)
    y = np.linspace(-Ny // 2, Ny // 2, Ny)
    X, Y = np.meshgrid(x, y)
    gaussian = amplitude * np.exp(-(((X - x0) ** 2) / (2 * sigma_x ** 2) +
                                    ((Y - y0) ** 2) / (2 * sigma_y ** 2)))
    return gaussian
In [7]:
ra_0, dec_0 = 0,0 # center of the frame
dx, dy = 1,  2 # we shift up and to the right these are in units of pixel/day for the projected velocities
MAG = 22.51 # magnitude of the injected star
In [8]:
total_data = {}
pass_counter = 0


for i in tqdm(range(len(files))):

    # Open the FITS file
    filename = files[i]
    hdul = fits.open(filename)    
    # Access the primary HDU (Header Data Unit)
    primary_hdu = hdul[0]
    
    # Get the data and header
    data = primary_hdu.data
    header = primary_hdu.header
    # ====== make sure to keep only images that are NOT all NANs
    total_elements = data.size
    # Count the number of NaN values
    nan_count = np.isnan(data).sum()
    # Calculate the percentage of NaN values
    nan_percentage = (nan_count / total_elements) * 100
    if nan_percentage < 20:
        time = header['SHUTOPEN']
        time = Time(time.replace('T',' '), format='iso', scale='utc')
        # ========== step 1: remove saturated pixels ==========
        saturate = header['SATURATE']
        data = np.where(data > saturate/2, np.nan, data) 
        # =========  step 2: subtract the median (sky subtraction) =======
        median = np.nanmedian(data)
        data = data - median
        # =========  step 3: scale flux =============
        MAGZP = header['MAGZP']
        scale = 10**(0.4*(MAGZPL - MAGZP))
        data = data * scale
        # =========== step 3.5: compute limiting Magnitude ======
        see = header['SEEING']
        pix = header['PIXSCALE']
        var = np.nanmedian(abs(data))
        lmt = -2.5*np.log10(3.0*np.sqrt(np.pi*see*see/pix/pix)*var*1.48) + MAGZPL
        # =========== step 4: drop gaussian fake =========
        grid_size = data.shape[::-1] # Grid size (Nx, Ny)
        dt = (time - earliest_frame).to_value('day') 
        x0, y0 = 0 + dt * dx, 0 + dt * dy  # shift the gaussian
        
        FWHM = header['SEEING']
        sigma = FWHM/2.3548
        sigma_x, sigma_y = sigma, sigma  # Standard deviations
        amplitude = 10**(-0.4*(MAG-MAGZPL)) *(1/(2*np.pi*sigma**2))
        gaussian = generate_2d_gaussian(grid_size, amplitude, x0, y0, sigma_x, sigma_y)
        print(gaussian.max())
        data = np.array(data) + gaussian
       
        header = primary_hdu.header
        
        total_data[time] = torch.tensor(data.astype(np.float32))
    else:
        pass_counter += 1
        pass
    hdul.close()
print(pass_counter, 'skipped out of ', len(files))  
# 8.940820307431759
  1%|▏                                          | 5/939 [00:00<01:09, 13.51it/s]
0.0
0.0
  1%|▎                                          | 7/939 [00:00<02:26,  6.36it/s]
8.85885037414935
  1%|▎                                          | 8/939 [00:01<02:51,  5.42it/s]
0.0
  1%|▍                                         | 10/939 [00:01<02:43,  5.68it/s]
8.469222277449635
  1%|▌                                         | 13/939 [00:01<02:16,  6.77it/s]
0.0
  3%|█▏                                        | 26/939 [00:02<00:52, 17.46it/s]
0.0
  3%|█▎                                        | 29/939 [00:02<01:02, 14.56it/s]
0.0
  4%|█▍                                        | 33/939 [00:03<01:07, 13.50it/s]
0.0
0.0
  4%|█▌                                        | 35/939 [00:03<01:44,  8.67it/s]
0.0
  4%|█▋                                        | 37/939 [00:03<01:51,  8.06it/s]
9.960558967402308
  4%|█▋                                        | 39/939 [00:04<01:59,  7.51it/s]
8.21863630660517
  5%|██                                        | 45/939 [00:04<01:26, 10.29it/s]
0.0
  5%|██                                        | 47/939 [00:04<01:36,  9.23it/s]
0.0
  6%|██▎                                       | 52/939 [00:05<01:23, 10.61it/s]
8.990446271057513
  6%|██▍                                       | 54/939 [00:05<01:32,  9.59it/s]
8.646129579492705
8.780614316596653
  6%|██▌                                       | 56/939 [00:06<02:06,  6.98it/s]
0.0
  6%|██▌                                       | 58/939 [00:06<02:04,  7.05it/s]
0.0
  6%|██▋                                       | 59/939 [00:06<02:21,  6.20it/s]
0.0
  7%|██▊                                       | 63/939 [00:07<01:50,  7.93it/s]
0.0
  7%|███                                       | 69/939 [00:07<01:21, 10.69it/s]
11.560324440670309
  8%|███▌                                      | 79/939 [00:07<00:48, 17.69it/s]
0.0
  9%|███▋                                      | 82/939 [00:08<00:57, 14.88it/s]
10.602419030858696
  9%|███▊                                      | 85/939 [00:08<01:05, 13.05it/s]
0.0
 10%|████                                      | 91/939 [00:08<00:59, 14.28it/s]
0.0
 10%|████▏                                     | 93/939 [00:09<01:10, 11.95it/s]
0.0
 10%|████▏                                     | 95/939 [00:09<01:21, 10.34it/s]
0.0
 10%|████▎                                     | 97/939 [00:09<01:33,  8.99it/s]
0.0
 11%|████▍                                     | 99/939 [00:10<01:40,  8.34it/s]
7.872451856853107
 11%|████▎                                    | 100/939 [00:10<02:00,  6.93it/s]
0.0
 11%|████▍                                    | 102/939 [00:10<02:03,  6.80it/s]
8.92602434329663
 12%|████▊                                    | 110/939 [00:11<01:11, 11.62it/s]
0.0
 12%|████▉                                    | 113/939 [00:11<01:15, 10.95it/s]
0.0
 13%|█████▏                                   | 120/939 [00:11<01:00, 13.59it/s]
7.99756145838178
 13%|█████▍                                   | 124/939 [00:12<01:02, 13.05it/s]
8.949795162106613
 14%|█████▋                                   | 129/939 [00:12<01:00, 13.31it/s]
3.537097260371463e-103
0.0
 14%|█████▋                                   | 131/939 [00:13<01:28,  9.11it/s]
11.626076068809411
 14%|█████▊                                   | 133/939 [00:13<01:35,  8.46it/s]
0.0
 14%|█████▊                                   | 134/939 [00:13<01:49,  7.37it/s]
9.429608273526904
 15%|██████                                   | 138/939 [00:14<01:33,  8.59it/s]
0.0
 15%|██████▏                                  | 141/939 [00:14<01:33,  8.55it/s]
0.0
 16%|██████▋                                  | 153/939 [00:14<00:45, 17.26it/s]
0.0
7.337030572541633
 17%|██████▊                                  | 157/939 [00:15<01:05, 11.90it/s]
8.8672213750851
0.0
 18%|███████▏                                 | 166/939 [00:16<00:58, 13.13it/s]
0.0
0.0
 18%|███████▍                                 | 169/939 [00:16<01:18,  9.75it/s]
0.0
 18%|███████▍                                 | 171/939 [00:17<01:26,  8.88it/s]
0.0
 19%|███████▋                                 | 176/939 [00:17<01:12, 10.46it/s]
0.0
 19%|███████▉                                 | 181/939 [00:17<01:05, 11.56it/s]
0.0
 19%|███████▉                                 | 183/939 [00:18<01:14, 10.13it/s]
0.0
 20%|████████                                 | 185/939 [00:18<01:23,  9.06it/s]
0.0
 20%|████████▏                                | 187/939 [00:18<01:29,  8.38it/s]
0.0
 20%|████████▏                                | 188/939 [00:19<01:47,  7.01it/s]
7.9864744829798315
 20%|████████▎                                | 189/939 [00:19<02:06,  5.92it/s]
7.1770921243383015
 21%|████████▍                                | 193/939 [00:19<01:35,  7.78it/s]
0.0
 21%|████████▌                                | 196/939 [00:19<01:30,  8.20it/s]
0.0
 22%|████████▊                                | 203/939 [00:20<01:03, 11.52it/s]
8.888872362326113
 22%|█████████                                | 207/939 [00:20<01:02, 11.72it/s]
0.0
 22%|█████████▏                               | 209/939 [00:21<01:12, 10.05it/s]
0.0
 23%|█████████▎                               | 214/939 [00:21<01:03, 11.35it/s]
0.0
 23%|█████████▍                               | 216/939 [00:21<01:12,  9.96it/s]
11.790833438229868
 23%|█████████▌                               | 218/939 [00:21<01:20,  8.98it/s]
0.0
 23%|█████████▌                               | 219/939 [00:22<01:39,  7.26it/s]
0.0
 23%|█████████▌                               | 220/939 [00:22<01:58,  6.08it/s]
8.680404807813357
 24%|█████████▋                               | 221/939 [00:22<02:15,  5.28it/s]
8.680824382848453
 24%|█████████▉                               | 227/939 [00:23<01:17,  9.14it/s]
0.0
 24%|█████████▉                               | 228/939 [00:23<01:33,  7.64it/s]
10.793525426413199
 25%|██████████                               | 231/939 [00:23<01:25,  8.24it/s]
0.0
 25%|██████████▏                              | 232/939 [00:24<01:42,  6.90it/s]
0.0
 25%|██████████▎                              | 235/939 [00:24<01:31,  7.69it/s]
0.0
 25%|██████████▎                              | 236/939 [00:24<01:48,  6.49it/s]
0.0
 26%|██████████▌                              | 241/939 [00:25<01:17,  9.06it/s]
0.0
 26%|██████████▊                              | 247/939 [00:25<00:59, 11.66it/s]
0.0
0.0
 27%|██████████▊                              | 249/939 [00:26<01:26,  7.95it/s]
0.0
 27%|██████████▉                              | 251/939 [00:26<01:30,  7.61it/s]
0.0
 27%|███████████                              | 252/939 [00:26<01:46,  6.45it/s]
0.0
 27%|███████████▏                             | 256/939 [00:27<01:25,  7.98it/s]
0.0
 28%|███████████▌                             | 265/939 [00:27<00:45, 14.91it/s]
3.638477679756935e-182
 30%|████████████▎                            | 282/939 [00:27<00:25, 25.82it/s]
9.244971292664893
 30%|████████████▍                            | 286/939 [00:28<00:32, 20.24it/s]
0.0
 31%|████████████▋                            | 291/939 [00:28<00:36, 17.86it/s]
0.0
0.0
 32%|█████████████▏                           | 302/939 [00:29<00:34, 18.51it/s]
10.913249026479741
 34%|█████████████▉                           | 320/939 [00:29<00:17, 35.65it/s]
0.0
 35%|██████████████▏                          | 326/939 [00:30<00:30, 20.08it/s]
9.625584097568575
 35%|██████████████▍                          | 331/939 [00:30<00:32, 18.54it/s]
0.0
 36%|██████████████▋                          | 335/939 [00:30<00:36, 16.44it/s]
0.0
7.7997662890398844
 36%|██████████████▊                          | 338/939 [00:31<00:53, 11.32it/s]
0.0
 36%|██████████████▉                          | 341/939 [00:31<00:55, 10.70it/s]
0.0
9.962447113294015
 37%|██████████████▉                          | 343/939 [00:32<01:15,  7.86it/s]
0.0
 37%|███████████████                          | 346/939 [00:32<01:12,  8.23it/s]
0.0
 37%|███████████████▏                         | 349/939 [00:33<01:08,  8.63it/s]
8.739870643397094
 38%|███████████████▌                         | 356/939 [00:33<00:49, 11.68it/s]
7.825880720286994
 38%|███████████████▋                         | 358/939 [00:33<00:56, 10.35it/s]
0.0
 39%|███████████████▉                         | 365/939 [00:34<00:44, 12.96it/s]
10.111796599506544
 39%|████████████████                         | 367/939 [00:34<00:50, 11.23it/s]
0.0
 40%|████████████████▏                        | 372/939 [00:34<00:46, 12.19it/s]
8.590996642247092
 40%|████████████████▎                        | 374/939 [00:35<00:53, 10.51it/s]
8.348012525958113
 40%|████████████████▍                        | 376/939 [00:35<01:01,  9.19it/s]
8.462572517428748
 40%|████████████████▍                        | 377/939 [00:35<01:13,  7.68it/s]
0.0
 41%|████████████████▋                        | 383/939 [00:36<00:52, 10.60it/s]
0.0
 41%|████████████████▊                        | 385/939 [00:36<00:59,  9.30it/s]
0.0
 41%|████████████████▉                        | 387/939 [00:36<01:05,  8.47it/s]
6.327993710476789
 42%|█████████████████                        | 390/939 [00:36<01:02,  8.74it/s]
10.037949575542566
 42%|█████████████████                        | 391/939 [00:37<01:15,  7.31it/s]
8.335274898305517
 42%|█████████████████▏                       | 395/939 [00:37<01:02,  8.75it/s]
0.0
 42%|█████████████████▎                       | 396/939 [00:37<01:15,  7.16it/s]
8.767673303095503
 42%|█████████████████▍                       | 398/939 [00:38<01:17,  6.94it/s]
0.0
 44%|█████████████████▊                       | 409/939 [00:38<00:31, 16.79it/s]
8.180215998234347
 44%|██████████████████                       | 413/939 [00:38<00:35, 14.86it/s]
0.0
 44%|██████████████████▏                      | 417/939 [00:39<00:37, 14.02it/s]
0.0
0.0
 45%|██████████████████▎                      | 420/939 [00:39<00:54,  9.58it/s]
0.0
 45%|██████████████████▍                      | 422/939 [00:40<00:58,  8.88it/s]
0.0
 46%|██████████████████▊                      | 432/939 [00:40<00:31, 16.12it/s]
0.0
 46%|███████████████████                      | 436/939 [00:40<00:34, 14.50it/s]
0.0
 47%|███████████████████▏                     | 439/939 [00:41<00:39, 12.82it/s]
8.263944618733122
 47%|███████████████████▎                     | 442/939 [00:41<00:42, 11.70it/s]
0.0
0.0
 47%|███████████████████▍                     | 444/939 [00:42<01:03,  7.77it/s]
8.146712263147615
 48%|███████████████████▊                     | 455/939 [00:42<00:31, 15.47it/s]
0.0
7.975878758837585
9.295504175092312
 49%|████████████████████                     | 459/939 [00:43<00:51,  9.26it/s]
0.0
 49%|████████████████████▏                    | 462/939 [00:43<00:51,  9.23it/s]
7.4171889686209
 50%|████████████████████▍                    | 469/939 [00:44<00:39, 11.92it/s]
0.0
 50%|████████████████████▌                    | 471/939 [00:44<00:43, 10.68it/s]
0.0
 51%|████████████████████▊                    | 476/939 [00:44<00:39, 11.67it/s]
8.04821653574666
 51%|████████████████████▊                    | 478/939 [00:45<00:44, 10.34it/s]
0.0
 51%|█████████████████████                    | 483/939 [00:45<00:40, 11.31it/s]
8.535200214225105
 52%|█████████████████████▏                   | 486/939 [00:45<00:41, 10.87it/s]
0.0
 52%|█████████████████████▍                   | 491/939 [00:46<00:37, 11.91it/s]
0.0
 53%|█████████████████████▌                   | 494/939 [00:46<00:40, 10.97it/s]
7.788103975565574
 53%|█████████████████████▋                   | 496/939 [00:46<00:44,  9.93it/s]
0.0
 53%|█████████████████████▋                   | 498/939 [00:47<00:50,  8.73it/s]
5.7835085156150345
 54%|██████████████████████▏                  | 507/939 [00:47<00:32, 13.47it/s]
0.0
 54%|██████████████████████▏                  | 509/939 [00:47<00:36, 11.66it/s]
1.80215e-319
 55%|██████████████████████▎                  | 512/939 [00:48<00:39, 10.81it/s]
0.0
 55%|██████████████████████▍                  | 514/939 [00:48<00:44,  9.51it/s]
10.466298845934949
 55%|██████████████████████▌                  | 516/939 [00:48<00:49,  8.61it/s]
0.0
 55%|██████████████████████▋                  | 519/939 [00:49<00:48,  8.75it/s]
0.0
 55%|██████████████████████▋                  | 520/939 [00:49<00:58,  7.17it/s]
0.0
 56%|██████████████████████▉                  | 526/939 [00:49<00:39, 10.40it/s]
0.0
 56%|███████████████████████                  | 528/939 [00:50<00:43,  9.44it/s]
0.0
 57%|███████████████████████▍                 | 536/939 [00:50<00:29, 13.68it/s]
0.0
 58%|███████████████████████▌                 | 541/939 [00:50<00:29, 13.65it/s]
10.160697358907935
 58%|███████████████████████▊                 | 544/939 [00:51<00:32, 12.25it/s]
0.0
8.895456451604382
 58%|███████████████████████▊                 | 546/939 [00:51<00:45,  8.58it/s]
9.441536611193635
 58%|███████████████████████▉                 | 547/939 [00:52<00:52,  7.43it/s]
0.0
 58%|███████████████████████▉                 | 549/939 [00:52<00:54,  7.12it/s]
9.195175892270328
 59%|████████████████████████                 | 551/939 [00:52<00:54,  7.10it/s]
0.0
 59%|████████████████████████▏                | 553/939 [00:52<00:55,  6.91it/s]
0.0
 59%|████████████████████████▏                | 555/939 [00:53<00:57,  6.72it/s]
0.0
 60%|████████████████████████▌                | 562/939 [00:53<00:35, 10.52it/s]
0.0
 60%|████████████████████████▋                | 566/939 [00:54<00:33, 10.99it/s]
0.0
 61%|█████████████████████████▏               | 576/939 [00:54<00:19, 18.33it/s]
0.0
0.0
 62%|█████████████████████████▎               | 580/939 [00:54<00:29, 12.29it/s]
0.0
 63%|█████████████████████████▋               | 587/939 [00:55<00:24, 14.44it/s]
8.54540892493153
 63%|█████████████████████████▊               | 591/939 [00:55<00:25, 13.84it/s]
0.0
 63%|█████████████████████████▉               | 593/939 [00:55<00:28, 11.97it/s]
0.0
 63%|█████████████████████████▉               | 595/939 [00:56<00:32, 10.55it/s]
0.0
 65%|██████████████████████████▍              | 606/939 [00:56<00:17, 18.55it/s]
0.0
10.597794679042611
 65%|██████████████████████████▋              | 610/939 [00:57<00:26, 12.37it/s]
0.0
 65%|██████████████████████████▊              | 613/939 [00:57<00:28, 11.42it/s]
0.0
 65%|██████████████████████████▊              | 615/939 [00:57<00:32, 10.06it/s]
0.0
 66%|██████████████████████████▉              | 617/939 [00:58<00:34,  9.30it/s]
0.0
 66%|███████████████████████████              | 620/939 [00:58<00:34,  9.38it/s]
0.0
 66%|███████████████████████████▏             | 622/939 [00:58<00:37,  8.45it/s]
9.426456963386864
 66%|███████████████████████████▏             | 624/939 [00:59<00:40,  7.81it/s]
11.163945850104055
 67%|███████████████████████████▎             | 625/939 [00:59<00:48,  6.52it/s]
0.0
 67%|███████████████████████████▍             | 627/939 [00:59<00:47,  6.55it/s]
4.7849948275249223e-14
 67%|███████████████████████████▌             | 632/939 [01:00<00:33,  9.05it/s]
0.0
 68%|███████████████████████████▋             | 634/939 [01:00<00:36,  8.29it/s]
0.0
 68%|███████████████████████████▊             | 637/939 [01:00<00:34,  8.73it/s]
0.0
 68%|███████████████████████████▉             | 641/939 [01:01<00:30,  9.79it/s]
0.0
 68%|████████████████████████████             | 642/939 [01:01<00:38,  7.81it/s]
0.0
 69%|████████████████████████████▏            | 646/939 [01:01<00:32,  9.02it/s]
7.945201652463195
 69%|████████████████████████████▎            | 647/939 [01:02<00:39,  7.36it/s]
9.410667766034113
 69%|████████████████████████████▎            | 649/939 [01:02<00:40,  7.12it/s]
0.0
 69%|████████████████████████████▍            | 652/939 [01:02<00:37,  7.74it/s]
0.0
 70%|████████████████████████████▋            | 658/939 [01:03<00:26, 10.63it/s]
9.824231569785058
 70%|████████████████████████████▊            | 661/939 [01:03<00:27, 10.07it/s]
0.0
 71%|████████████████████████████▉            | 663/939 [01:03<00:30,  8.97it/s]
0.0
 71%|█████████████████████████████            | 665/939 [01:03<00:32,  8.52it/s]
0.0
 72%|█████████████████████████████▌           | 678/939 [01:04<00:13, 18.90it/s]
0.0
 73%|█████████████████████████████▊           | 682/939 [01:04<00:15, 16.48it/s]
0.0
 73%|█████████████████████████████▉           | 685/939 [01:05<00:17, 14.34it/s]
0.0
0.0
 73%|██████████████████████████████           | 688/939 [01:05<00:25,  9.96it/s]
8.006180442862936
 74%|██████████████████████████████▎          | 695/939 [01:06<00:19, 12.60it/s]
0.0
 74%|██████████████████████████████▍          | 698/939 [01:06<00:20, 11.53it/s]
10.848144347715635
 75%|██████████████████████████████▋          | 703/939 [01:06<00:18, 12.60it/s]
0.0
 75%|██████████████████████████████▊          | 705/939 [01:07<00:21, 11.01it/s]
0.0
11.14559498495139
 75%|██████████████████████████████▊          | 707/939 [01:07<00:30,  7.71it/s]
0.0
 75%|██████████████████████████████▉          | 708/939 [01:07<00:34,  6.71it/s]
0.0
 76%|██████████████████████████████▉          | 709/939 [01:08<00:39,  5.79it/s]
0.0
 76%|███████████████████████████████          | 711/939 [01:08<00:38,  5.95it/s]
8.461355962756768
 77%|███████████████████████████████▍         | 719/939 [01:08<00:20, 10.63it/s]
8.065837104722625
 77%|███████████████████████████████▋         | 727/939 [01:09<00:15, 13.70it/s]
0.0
 78%|███████████████████████████████▊         | 729/939 [01:09<00:17, 11.83it/s]
0.0
 79%|████████████████████████████████▎        | 739/939 [01:10<00:11, 18.13it/s]
0.0
0.0
 79%|████████████████████████████████▍        | 743/939 [01:10<00:15, 12.31it/s]
0.0
 79%|████████████████████████████████▌        | 746/939 [01:10<00:16, 11.40it/s]
9.781389492706348
 80%|████████████████████████████████▋        | 749/939 [01:11<00:17, 10.74it/s]
0.0
 80%|████████████████████████████████▊        | 752/939 [01:11<00:18, 10.23it/s]
9.335522609032125
 82%|█████████████████████████████████▍       | 766/939 [01:12<00:08, 19.66it/s]
0.0
11.36801469157964
 82%|█████████████████████████████████▌       | 770/939 [01:12<00:12, 13.41it/s]
0.0
0.0
 82%|█████████████████████████████████▊       | 773/939 [01:13<00:16, 10.05it/s]
0.0
 83%|█████████████████████████████████▊       | 775/939 [01:13<00:18,  9.07it/s]
0.0
 83%|█████████████████████████████████▉       | 777/939 [01:13<00:19,  8.49it/s]
0.0
 83%|██████████████████████████████████       | 779/939 [01:14<00:20,  7.83it/s]
0.0
 84%|██████████████████████████████████▎      | 785/939 [01:14<00:14, 10.39it/s]
0.0
 84%|██████████████████████████████████▍      | 788/939 [01:14<00:15,  9.88it/s]
0.0
 85%|██████████████████████████████████▊      | 796/939 [01:15<00:10, 13.58it/s]
0.0
 85%|██████████████████████████████████▉      | 799/939 [01:15<00:11, 12.43it/s]
7.242386051038754
 86%|███████████████████████████████████▏     | 805/939 [01:16<00:09, 13.66it/s]
7.859979486506908
0.0
 87%|███████████████████████████████████▌     | 815/939 [01:16<00:07, 15.74it/s]
10.544880749031611
7.489815460847917
8.439965735485005
 87%|███████████████████████████████████▊     | 819/939 [01:17<00:12,  9.52it/s]
0.0
 88%|███████████████████████████████████▉     | 824/939 [01:17<00:10, 10.64it/s]
7.885413581569076
 88%|████████████████████████████████████▏    | 830/939 [01:18<00:09, 12.03it/s]
0.0
 89%|████████████████████████████████████▎    | 832/939 [01:18<00:09, 10.81it/s]
0.0
 89%|████████████████████████████████████▌    | 836/939 [01:19<00:09, 10.96it/s]
10.438028695114324
 90%|████████████████████████████████████▊    | 844/939 [01:19<00:06, 14.07it/s]
9.372966498151573
 90%|████████████████████████████████████▉    | 846/939 [01:19<00:07, 12.01it/s]
8.913214550392658
0.0
 90%|█████████████████████████████████████    | 848/939 [01:20<00:10,  8.30it/s]
7.248676652190477
 91%|█████████████████████████████████████▏   | 851/939 [01:20<00:10,  8.45it/s]
0.0
 91%|█████████████████████████████████████▍   | 858/939 [01:21<00:07, 11.38it/s]
9.817345813124907
 92%|█████████████████████████████████████▌   | 861/939 [01:21<00:07, 10.91it/s]
0.0
 92%|█████████████████████████████████████▋   | 863/939 [01:21<00:07,  9.60it/s]
8.940564423926785
 92%|█████████████████████████████████████▊   | 867/939 [01:21<00:06, 10.47it/s]
0.0
 94%|██████████████████████████████████████▍  | 879/939 [01:22<00:03, 18.79it/s]
0.0
 94%|██████████████████████████████████████▋  | 887/939 [01:22<00:02, 19.25it/s]
8.995597144251777
 96%|███████████████████████████████████████▎ | 899/939 [01:23<00:01, 25.16it/s]
8.039929379091456
0.0
 96%|███████████████████████████████████████▍ | 904/939 [01:23<00:02, 16.64it/s]
0.7698665672940787
 97%|███████████████████████████████████████▋ | 908/939 [01:24<00:02, 15.12it/s]
7.8854746272599385
 97%|███████████████████████████████████████▊ | 913/939 [01:24<00:01, 14.84it/s]
0.0
0.0
 98%|███████████████████████████████████████▉ | 916/939 [01:25<00:02, 10.50it/s]
9.40167462182834
0.0
 98%|████████████████████████████████████████ | 918/939 [01:25<00:02,  7.83it/s]
10.678290666912067
 98%|████████████████████████████████████████▏| 921/939 [01:26<00:02,  8.40it/s]
0.0
 99%|████████████████████████████████████████▍| 925/939 [01:26<00:01,  9.27it/s]
7.875649136321271
7.719297637567392
 99%|████████████████████████████████████████▍| 927/939 [01:26<00:01,  6.97it/s]
0.0
 99%|████████████████████████████████████████▋| 931/939 [01:27<00:00,  8.13it/s]
9.076783379425256
100%|████████████████████████████████████████▊| 935/939 [01:27<00:00,  8.99it/s]
0.0
100%|████████████████████████████████████████▉| 938/939 [01:27<00:00,  9.09it/s]
0.0
100%|█████████████████████████████████████████| 939/939 [01:28<00:00, 10.64it/s]
0.0
671 skipped out of  939

In [9]:
sorted_keys = sorted(total_data.keys())
earliest_frame = sorted_keys[0]

Watch the injections!¶

Let us visualise the fake injected star

In [10]:
from astropy.io import fits
import matplotlib.pyplot as plt
from astropy.visualization import (ImageNormalize, ZScaleInterval, LogStretch, MinMaxInterval, AsinhStretch)
import numpy as np
import matplotlib.pyplot as plt
from celluloid import Camera
from IPython.display import HTML

norm = ImageNormalize(stretch=AsinhStretch(a=0.0001), interval=ZScaleInterval())  

%matplotlib inline
# Set up the figure and axes
fig, ax = plt.subplots(figsize=(10,10))

# Initialize Celluloid Camera
camera = Camera(fig)


data = total_data[earliest_frame]
count = 0
# Animate the sine wave with a changing phase
for key in sorted_keys:
    ax.imshow(total_data[key], cmap='gray_r', origin='lower', norm=norm)  # Plot the sine wave
    camera.snap()  # Capture the frame
    if count > 100:
        break
    else:
        count += 1

# Create the animation
animation = camera.animate(interval=100)  # 100ms between frames

# Display the animation
HTML(animation.to_jshtml())  # Render as JS HTML
Out[10]:
No description has been provided for this image
No description has been provided for this image
In [11]:
import torch
import torch.nn as nn
import torch.nn.functional as F
from time import time

class Shift_Stack(nn.Module):
    def __init__(self, frames, times):
        super(Shift_Stack, self).__init__()
        self.frames = self.frame_subtraction(frames).cuda()[:,None,:,:]
        self.delta_times = torch.zeros_like(times).cuda()
        self.delta_times[1:] = times[1:] - times[0]
        self.B = self.frames.shape[0] # predefine batch sizes (number of frames)
        self.grid = F.affine_grid(
            torch.eye(2, 3).unsqueeze(0).repeat(self.B, 1, 1),  # Identity matrix for rotation and scaling
            size=self.frames.size(),  # Output size
            align_corners=False
        ).cuda()
        
    # @torch.jit.script
    def frame_subtraction(self, frame):
        """
        Sequential frame subtractions
        """
        frame = frame - frame[0,:,:]
        return frame
        
    def decimate_32_to_16(self, data_float32):
        """
        Decimate data from float32 to float16 while preserving the range
        """
        # Define a scaling factor to bring values within float16 range
        scale_factor = data_float32.abs().max() / 6.55e4  # Approx. max for float16
        scale_factor = max(scale_factor, 1.0)  # Ensure we only scale down
        
        # Scale, convert to float16, and scale back
        data_scaled = data_float32 / scale_factor
        data_float16 = data_scaled.to(dtype=torch.float16)
        data_preserved = data_float16 * scale_factor
        return data_preserved
        
    def average_nonzeros(self, data, axis=0):
        """
        This averages only the nonzero values
        """
        # Mask to identify non-zero values
        non_zero_sum = torch.sum(data, dim=axis)
        non_zero_count = torch.count_nonzero(data, dim=axis)
        
        # Avoid division by zero
        non_zero_count = torch.where(non_zero_count == 0, torch.tensor(1.0, device=data.device), non_zero_count)
        
        # Calculate average of non-zero elements
        non_zero_avg = non_zero_sum / non_zero_count
        return non_zero_avg

    def shift_images(self, images, shifts):
        """
        Shifts a batch of images by specified (x, y) values for each image.
        
        images (torch.Tensor): A batch of images of shape (B, C, H, W),
                               where B is batch size, C is number of channels,
                               H is height, and W is width.
        shifts (torch.Tensor): A tensor of shape (B, 2), where each row
                               represents the (B, x, y) shift for the corresponding image.
        """
        B, C, H, W = images.shape
        shifts_x, shifts_y = shifts[:, 0], shifts[:, 1]
        # Create a grid for affine transformations
        start = time()
        grid = self.grid.clone()
        # Adjust the grid by the shifts
        grid[..., 0] += 2 * shifts_x.view(-1, 1, 1) / W  # Normalize shift in x direction
        grid[..., 1] += 2 * shifts_y.view(-1, 1, 1) / H  # Normalize shift in y direction
        start = time()
        # Perform the grid sampling to shift the images
        shifted_images = F.grid_sample(
            images, grid, mode='bilinear', padding_mode='zeros', align_corners=False
        ).cuda()
        return shifted_images

    def forward(self, dxdy):
        """
        Preprocessing handing before calling shift_images
        """
        dx, dy = dxdy[:,0], dxdy[:,1]
        shifts_x = torch.multiply(dx , self.delta_times) # pixel shift in x
        shifts_y = torch.multiply(dy , self.delta_times) # pixel shift in y
        shifts = torch.stack([shifts_x, shifts_y] , axis = 1)
        frames = self.frames
        batched_shift = self.shift_images(frames, shifts)
        batched_shift = batched_shift[:,0,:,: ] # meant for future multi-filters... [curr. index zero for 1 filter]
        # stacked = torch.mean(batched_shift, axis = 0) # FASTER MEAN Calculations
        stacked = torch.nanmedian(batched_shift, axis = 0)
        # stacked = self.average_nonzeros(batched_shift, axis = 0)
        return stacked

Shift and Add¶

We write the batched version of shift and add model

Load up the data into torch tensors.

In [12]:
frames = []
frame_times = []
adjusts = []
for i, key in enumerate(sorted_keys):
    if i < 150:
        frames.append(total_data[key])
        frame_times.append(key.to_value('jd'))
frames = torch.stack(frames)
frame_times = torch.tensor(frame_times)
print(frames.shape, frame_times.shape)
torch.Size([150, 3080, 3072]) torch.Size([150])

Instantiate a SINGLE shift stack model for 1 single frame

In [13]:
stacking_job = Shift_Stack(frames, frame_times)

Run a job on shift and stacking for 1 sequence of orbits

In [14]:
shifts = torch.ones((150, 2)).cuda()
shifts[:, 0] = dx
shifts[:, 1] = dy
print(shifts.shape)
torch.Size([150, 2])

Shifts data object is in the following shape $[time, spatial]$. Here there are 100 frames at different times, each frame has a corresponding 2d adjustment vector. For us since the velocity is constant in time we do not need to worry about different dx dy as a function of time.

In [15]:
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
start = time()
result = stacking_job(shifts)
print(time()-start, "runtime [s]")
0.12030720710754395 runtime [s]
0.0007495880126953125 runtime [s]
0.00016570091247558594 runtime [s]
0.00015783309936523438 runtime [s]
0.00015497207641601562 runtime [s]
0.0001575946807861328 runtime [s]
0.0001533031463623047 runtime [s]
In [22]:
norm = ImageNormalize(stretch=AsinhStretch(a=100), interval=ZScaleInterval())  
# 
plt.figure(figsize=(10, 10))
plt.imshow(result.values.cpu().numpy(), cmap='gray_r', origin='lower', vmin = -10, vmax=10 )
# plt.imshow(result.values.cpu().numpy(), cmap='gray_r', origin='lower', vmin = result.values.cpu().numpy().min(), vmax=result.values.cpu().numpy().max() )
plt.colorbar(label='Pixel Value')
plt.title('FITS Image with Normalization')
plt.xlabel('RA Pixel')
plt.ylabel('Dec Pixel')
plt.show()
No description has been provided for this image
In [23]:
plt.figure(figsize=(10, 10))
plt.imshow(result.values.cpu().numpy()[1500:1600,1500:1600], cmap='gray_r', origin='lower', vmin = result.values.cpu().numpy()[1500:1600,1500:1600].min(), vmax=result.values.cpu().numpy()[1500:1600,1500:1600].max() )
plt.colorbar(label='Pixel Value')
plt.title('FITS Image with Normalization')
plt.xlabel('RA Pixel')
plt.ylabel('Dec Pixel')
plt.show()
No description has been provided for this image
In [20]:
%matplotlib inline
In [21]:
array = result.values.cpu().numpy()
flat_index = np.argmax(array)

# Convert flat index to 2D indices
row, col = np.unravel_index(flat_index, data.shape)

print("Maximum value:", array[row, col])
print("Index of maximum value:", (row, col))
Maximum value: 10.842533
Index of maximum value: (np.int64(62), np.int64(2148))